/*******************************************************************************
 * CopyRight (c) 2020 
 *
 * \file    mkv.c
 * \brief   MKV file format parser
 * \author  ZhangJing
 * \version 1.0.0
 * \date    2020-06-22
 
 * -----------------------------------------------------------------------------
 * Modify history：
 * <Date>       | <Version> | <Author>  | <Discribe>
 * 2020-06-22   | v1.0.0    | ZhangJing | Create file
 * -----------------------------------------------------------------------------
 ******************************************************************************/

#include <stdio.h>
#include <string.h>
#include <stdlib.h>

#include "publicdef.h"
#include "print_ctrl.h"
#include "format.h"
#include "mkv.h"

//FIXME:
static FILE *oafp = NULL;
static FILE *ovfp = NULL;

/*!
 * \brief   mkv file format probe
 *
 * \param[in] format_context file data read
 * \return    0 if succeeded otherwise failed
 */
static int mkv_probe(void *format_context)
{
    FormatContext *fc = format_context;
    int ret = 0;
    unsigned int id;

    pr_entry(in);

    /** read EBML header */
    memcpy(&id, fc->probe_data, sizeof(id));

    /* match EBML Header id */
    if (id != MKLE32(EBML_ID_HEADER))
        return -EPROTONOSUPPORT;

    //TODO: check DocType = “matroska”

    return ret;
}

const unsigned char log2_tab[256]={
        0,0,1,1,2,2,2,2,3,3,3,3,3,3,3,3,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,
        5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,
        6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,
        6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,
        7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,
        7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,
        7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,
        7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7
};

/*!
 * \brief get real data from variable length data
 */
static int ebml_read_num(void *fm, int max_size, unsigned long *number)
{
    FormatContext *fc = fm;
    IOContext *io = fc->io;
    unsigned long total = 0;
    int read, n = 1;
    unsigned char tmp;

    /* The first byte tells us the length in bytes - except when it is zero. */
    io->io_read(io, &total, 1);

    /* get the length of the EBML number */
    read = 8 - log2_tab[total];

    /* read out length */
    total ^= 1 << log2_tab[total];
    while (n++ < read) {
        io->io_read(io, &tmp, 1);
        total = (total << 8) | tmp;
    }

    *number = total;

    return read;
}

/*!
 *\brief get embl id
 */
static MkvSyntax *ebml_parse_id(MkvSyntax *syntax, unsigned int id)
{
    int i;
static MkvSyntax null_syntax = {0, NULL, EBML_NONE, 0, 0};

    for (i = 0; syntax[i].id; i++) {
        if (id == syntax[i].id) {
            return &syntax[i];
        }
    }

    return &null_syntax;
}

/*!
 * \brief read trak element and alloc stream struct
 */
static int mkv_read_codecid(void *fm, MkvSyntax *syntax)
{
    FormatContext *fc = fm;
    IOContext *io = fc->io;
    Stream *st;
    Stream **streams;
    MkvStream *ms;
    char *str = NULL;

    if (fc->state == 1) { /**< if is write packet, skip */
        /* skip the data */
        io->io_seek(io, syntax->size, SEEK_CUR);
        return 0;
    }

    /** malloc stream */
    if ((fc->nb_streams <= 0) || !fc->streams) {
        streams = malloc(sizeof(*streams));
        if (!streams)
            return -ENOMEM;
    } else {
        streams = fc->streams;
        streams = realloc(streams, (fc->nb_streams + 1) * sizeof(*streams));
        if (!streams)
            return -ENOMEM;
    }
    fc->streams = streams;

    st = malloc(sizeof(Stream));
    if (!st)
        return -ENOMEM;
    st->index = fc->nb_streams;
    fc->streams[fc->nb_streams++] = st;

    ms = malloc(sizeof(MkvStream));
    if (!ms)
        return -ENOMEM;
    st->mkvst = ms;

    str = malloc(syntax->size + 1);
    if (!str) {
        /* skip the data */
        io->io_seek(io, syntax->size, SEEK_CUR);
        return -ENOMEM;
    }

    io->io_read(io, str, syntax->size);
    str[syntax->size] = '\0';

    pr_info("MKV Codec ID: %s", str);
    if (strstr(str, "AVC")) {
        st->streamname = "h264";
        st->id = ID_H264;
        st->track_type = 1; //FIXME: get this in embl TrackType(cmd 0x83)
    } else if (strstr(str, "AAC")) {
        st->streamname = "aac";
        st->id = ID_AAC;
        st->track_type = 2;
    } else if (strstr(str, "AC3")) {
        st->streamname = "aac";
        st->id = ID_AAC;
        st->track_type = 2;
    } else { //FIXME: //TODO: add others
        st->streamname = "aac";
        st->id = ID_AAC;
        st->track_type = 2;
    }

    if (st->streamname)
        pr_info("MKV stream name: %s", st->streamname);

    return 0;
}

#define BLOCK_BUF_SIZE (1024 * 1024 *20)
char block_buf[BLOCK_BUF_SIZE];
char obuf[BLOCK_BUF_SIZE];
/*!
 * \brief read packet block data
 */
static int mkv_read_block(void *fm, MkvSyntax *syntax)
{
    FormatContext *fc = fm;
    IOContext *io = fc->io;
    Stream *ost = fc->streams[0]; /**< if output */
    unsigned int block_header;
    unsigned char header_track;
    int track_av_type; /* 1 video, 2 audio */
    unsigned short timecode;
    unsigned char flag;

    if (fc->state == 0) { /**< if is read header info, skip */
        /* skip the data */
        io->io_seek(io, syntax->size, SEEK_CUR);
        return 0;
    }

    /** get video or audio type */
    io->io_read(io, &header_track, sizeof(header_track));
    //FIXME: use ebml_read_num() matching the track number
    track_av_type = (header_track == 0x81 ? 1 : 2);
    io->io_read(io, &timecode, sizeof(timecode));
    io->io_read(io, &flag, sizeof(flag));

    int olen = 0;
    if (track_av_type != ost->track_type) {
        /* skip the data */
        io->io_seek(io, syntax->size - sizeof(block_header), SEEK_CUR);
        return 0;
    }

    if (track_av_type == 1) {
        /* video, default h264*/
        io->io_read(io, block_buf, syntax->size - sizeof(block_header));
        if (ost->bsf) {
            ost->bsf->filter(fc, block_buf, syntax->size - sizeof(block_header), obuf, &olen);
            fwrite(obuf, 1, olen, ovfp);
        } else
            fwrite(block_buf, 1, syntax->size - sizeof(block_header), ovfp);
    } else {
        /* audio, default aac*/
        io->io_read(io, block_buf, syntax->size - sizeof(block_header));
        if (ost->bsf) {
            ost->bsf->filter(fc, block_buf, syntax->size - sizeof(block_header), obuf, &olen);
            fwrite(obuf, 1, olen, oafp);
        } else
            fwrite(block_buf, 1, syntax->size - sizeof(block_header), oafp);
    }

    return 0;
}

/*!
 * \brief read parent ebml and child ebml with recursion
 */
static int ebml_parse(void *fm, MkvSyntax *syntax)
{
    FormatContext *fc = fm;
    IOContext *io = fc->io;
    int res;
    unsigned long id;
    unsigned int current_id;
    unsigned long length;
    MkvSyntax *id_syntax;
static int count_test = 0;
    //FIXME: use array
    MkvSyntax elem;

    long seek_start;

    pr_detail_pure("-->\n");
    seek_start = io->io_tell(io);

    /** get id */
    res = ebml_read_num(fc, 4, &id);
    current_id = id | 1 << 7 * res;

    /** get option of the id */
    id_syntax = ebml_parse_id(syntax, current_id);

    if (id_syntax->id == EBML_ID_HEADER)
        count_test = 0;
    else
        count_test++;
    
    /** get length */
    res = ebml_read_num(fc, 8, &length);
    if (id_syntax->name) {
        pr_detail("id: <{%s}> (0x%x), len: 0x%lx, i: %d", id_syntax->name,
                 id_syntax->id, length, count_test);
    } else
        pr_info("!ERROR id: 0x%x, len: 0x%lx", current_id, length);

    if (id_syntax->id) {
        //FIXME: save the value in array
        elem.id = id_syntax->id;
        elem.name = id_syntax->name;
        elem.size = length;
        elem.offset = io->io_tell(io);
        elem.type = id_syntax->type;
    }
    if (id_syntax->parser)
        id_syntax->parser(fm, &elem);

    switch (id_syntax->type) {
    case EBML_BIN:
    case EBML_STRING:
        /* read bin data in parser function */
        break;
    case EBML_NONE:
    case EBML_SKIP:
        /* skip the data */
        io->io_seek(io, length, SEEK_CUR);
        break;
    case EBML_NEST:
        /** recursive call */
        pr_detail("seek_start: 0x%lx, seek_end:0x%lx", seek_start, io->io_tell(io));
        ebml_parse(fm, syntax);
        break;
    default:
        /* skip the data */
        io->io_seek(io, length, SEEK_CUR);
        break;
    }

    if (id_syntax->type != EBML_NEST) {
        pr_detail("seek_start: 0x%lx, seek_end:0x%lx", seek_start, io->io_tell(io));
    }

    return 0;
}

/*!
 * \brief get pps data
 */
static int mkv_read_private(void *fm, MkvSyntax *syntax)
{
    FormatContext *fc = fm;
    IOContext *io = fc->io;
    Stream *st;
    MkvStream *ms;
    int total_len = syntax->size;
    int cur_len = 0;
    unsigned char tag_size;
    unsigned int start_code4 = 0x01000000; /**< start code of h264 NALU */
    char start_code3[3] = {0x00, 0x00, 0x01};

    if (fc->state != 0) { /**< if is read packet, skip */
        /* skip the data */
        io->io_seek(io, syntax->size, SEEK_CUR);
        return 0;
    }

    if (fc->nb_streams < 1)
        goto exit;
    st = fc->streams[fc->nb_streams - 1];
    ms = st->mkvst;

    if (syntax->size < 25) //FIXME:
        goto exit;

    /* skip 5 bytes */
    io->io_seek(io, 5, SEEK_CUR); cur_len += 5;
    /* skip 2 bytes */
    io->io_seek(io, 2, SEEK_CUR); cur_len += 2;
    /* get pps len */
    io->io_read(io, &tag_size, 1); cur_len += 1;

    /* malloc pps data */
    ms->pps_data = malloc(128);
    if (!ms->pps_data)
        goto exit;

    ms->pps_data_len = 0;

    pr_info("1 pps len: %d", tag_size);
    memcpy(ms->pps_data, &start_code3, 3); ms->pps_data_len += 3;
    io->io_read(io, ms->pps_data + ms->pps_data_len, tag_size); cur_len += tag_size; ms->pps_data_len += tag_size;

    /* skip 2 bytes */
    io->io_seek(io, 2, SEEK_CUR); cur_len += 2;
    /* get pps len */
    io->io_read(io, &tag_size, 1); cur_len += 1;
    pr_info("2 pps len: %d", tag_size);
    memcpy(ms->pps_data + ms->pps_data_len, &start_code4, 4); ms->pps_data_len += 4;
    io->io_read(io, ms->pps_data + ms->pps_data_len, tag_size); cur_len += tag_size; ms->pps_data_len += tag_size;

exit:
    /* skip the bytes remaining */
    io->io_seek(io, total_len - cur_len, SEEK_CUR); cur_len += 4;

    return 0;
}

MkvSyntax mkv_syntax [] = {
    {EBML_ID_HEADER, name2str(EBML_ID_HEADER), EBML_NEST},
    {MKV_ID_SEGMENT, name2str(MKV_ID_SEGMENT), EBML_NEST},
    {MKV_ID_SEEKHEAD, name2str(MKV_ID_SEEKHEAD), EBML_NEST},
    {MKV_ID_SEEKENTRY, name2str(MKV_ID_SEEKENTRY), EBML_NEST},
    {MKV_ID_INFO, name2str(MKV_ID_INFO), EBML_NEST},
    {MKV_ID_TRACKS, name2str(MKV_ID_TRACKS), EBML_NEST},
    {MKV_ID_TRACKENTRY, name2str(MKV_ID_TRACKENTRY), EBML_NEST},
    {MKV_ID_TRACKVIDEO, name2str(MKV_ID_TRACKVIDEO), EBML_NEST},
    {MKV_ID_TRACKAUDIO, name2str(MKV_ID_TRACKAUDIO), EBML_NEST},
    {MKV_ID_CUES, name2str(MKV_ID_CUES), EBML_NEST},
    {MKV_ID_POINTENTRY, name2str(MKV_ID_POINTENTRY), EBML_NEST},
    {MKV_ID_CUETRACKPOSITION, name2str(MKV_ID_CUETRACKPOSITION), EBML_NEST},
    {MKV_ID_TAGS, name2str(MKV_ID_TAGS), EBML_NEST},
    {MKV_ID_TAG, name2str(MKV_ID_TAG), EBML_NEST},
    {MKV_ID_SIMPLETAG, name2str(MKV_ID_SIMPLETAG), EBML_NEST},
    {MKV_ID_CLUSTER, name2str(MKV_ID_CLUSTER), EBML_NEST},
    {EBML_ID_VOID, name2str(EBML_ID_VOID), EBML_SKIP},
    {EBML_ID_CRC32, name2str(EBML_ID_CRC32), EBML_SKIP},
    {EBML_ID_UNKNOWN, name2str(EBML_ID_UNKNOWN), EBML_SKIP},
    {EBML_ID_EBMLVERSION, name2str(EBML_ID_EBMLVERSION), EBML_SKIP},
    {EBML_ID_EBMLREADVERSION, name2str(EBML_ID_EBMLREADVERSION), EBML_SKIP},
    {EBML_ID_EBMLMAXIDLENGTH, name2str(EBML_ID_EBMLMAXIDLENGTH), EBML_SKIP},
    {EBML_ID_EBMLMAXSIZELENGTH, name2str(EBML_ID_EBMLMAXSIZELENGTH), EBML_SKIP},
    {EBML_ID_DOCTYPE, name2str(EBML_ID_DOCTYPE), EBML_SKIP},
    {EBML_ID_DOCTYPEVERSION, name2str(EBML_ID_DOCTYPEVERSION), EBML_SKIP},
    {EBML_ID_DOCTYPEREADVERSION, name2str(EBML_ID_DOCTYPEREADVERSION), EBML_SKIP},
    {MKV_ID_SEEKID, name2str(MKV_ID_SEEKID), EBML_SKIP},
    {MKV_ID_SEEKPOSITION, name2str(MKV_ID_SEEKPOSITION), EBML_SKIP},
    {MKV_ID_TIMECODESCALE, name2str(MKV_ID_TIMECODESCALE), EBML_SKIP},
    {MKV_ID_DURATION, name2str(MKV_ID_DURATION), EBML_SKIP},
    {MKV_ID_WRITINGAPP, name2str(MKV_ID_WRITINGAPP), EBML_SKIP},
    {MKV_ID_MUXINGAPP, name2str(MKV_ID_MUXINGAPP), EBML_SKIP},
    {MKV_ID_DATEUTC, name2str(MKV_ID_DATEUTC), EBML_SKIP},
    {MKV_ID_SEGMENTUID, name2str(MKV_ID_SEGMENTUID), EBML_SKIP},
    {MKV_ID_TRACKNUMBER, name2str(MKV_ID_TRACKNUMBER), EBML_SKIP},
    {MKV_ID_TRACKUID, name2str(MKV_ID_TRACKUID), EBML_SKIP},
    {MKV_ID_TRACKTYPE, name2str(MKV_ID_TRACKTYPE), EBML_SKIP}, //TODO: add
    {MKV_ID_VIDEODISPLAYWIDTH, name2str(MKV_ID_VIDEODISPLAYWIDTH), EBML_SKIP},
    {MKV_ID_VIDEODISPLAYHEIGHT, name2str(MKV_ID_VIDEODISPLAYHEIGHT), EBML_SKIP},
    {MKV_ID_VIDEOPIXELWIDTH, name2str(MKV_ID_VIDEOPIXELWIDTH), EBML_SKIP},
    {MKV_ID_VIDEOPIXELHEIGHT, name2str(MKV_ID_VIDEOPIXELHEIGHT), EBML_SKIP},
    {MKV_ID_VIDEOFLAGINTERLACED, name2str(MKV_ID_VIDEOFLAGINTERLACED), EBML_SKIP},
    {MKV_ID_AUDIOSAMPLINGFREQ, name2str(MKV_ID_AUDIOSAMPLINGFREQ), EBML_SKIP},
    {MKV_ID_AUDIOBITDEPTH, name2str(MKV_ID_AUDIOBITDEPTH), EBML_SKIP},
    {MKV_ID_AUDIOCHANNELS, name2str(MKV_ID_AUDIOCHANNELS), EBML_SKIP},
    {MKV_ID_CODECID, name2str(MKV_ID_CODECID), EBML_STRING, 0, 0, mkv_read_codecid},
    {MKV_ID_CODECPRIVATE, name2str(MKV_ID_CODECPRIVATE), EBML_BIN, 0, 0, mkv_read_private},
    {MKV_ID_TRACKLANGUAGE, name2str(MKV_ID_TRACKLANGUAGE), EBML_SKIP},
    {MKV_ID_TRACKFLAGLACING, name2str(MKV_ID_TRACKFLAGLACING), EBML_SKIP},
    {MKV_ID_TRACKMINCACHE, name2str(MKV_ID_TRACKMINCACHE), EBML_SKIP},
    {MKV_ID_TRACKDEFAULTDURATION, name2str(MKV_ID_TRACKDEFAULTDURATION), EBML_SKIP},
    {MKV_ID_CUETIME, name2str(MKV_ID_CUETIME), EBML_SKIP},
    {MKV_ID_CUECLUSTERPOSITION, name2str(MKV_ID_CUECLUSTERPOSITION), EBML_SKIP},
    {MKV_ID_TAGTARGETS, name2str(MKV_ID_TAGTARGETS), EBML_SKIP},
    {MKV_ID_TAGNAME, name2str(MKV_ID_TAGNAME), EBML_SKIP},
    {MKV_ID_TAGSTRING, name2str(MKV_ID_TAGSTRING), EBML_SKIP},
    {MKV_ID_CLUSTERTIMECODE, name2str(MKV_ID_CLUSTERTIMECODE), EBML_SKIP},
    {MKV_ID_SIMPLEBLOCK, name2str(MKV_ID_SIMPLEBLOCK), EBML_BIN, 0, 0, mkv_read_block},
    {MKV_ID_CUETRACK, name2str(MKV_ID_CUETRACK), EBML_SKIP},
    {MKV_ID_CUERELATIVEPOSITION, name2str(MKV_ID_CUERELATIVEPOSITION), EBML_SKIP},
    { 0 }
};

/*!
 * \brief     get stream information
 *
 * \param[in] format_context file data read
 * \return    0 if succeeded otherwise failed
 */
static int mkv_read_header(void *format_context)
{
    FormatContext *fc = format_context;
    int offset = 0;
    IOContext *io = fc->io;
    long file_len;
    
    pr_entry(in);

    file_len = io->io_get_len(io);
    io->io_seek(io, 0, SEEK_SET);

    fc->nb_streams = 0;
    fc->ofname = NULL;
    fc->ofname_audio = NULL;
    fc->state = 0; /**< not write packet */

    /* read all EBML elements. */
    do {
       ebml_parse(fc, mkv_syntax);
    } while (io->io_tell(io) < file_len);

    return 0;
}

/*!
 * \brief     get stream data
 *
 * \param[in] format_context file data read
 * \return    0 if succeeded otherwise failed
 */
static int mkv_read_packet(void *format_context)
{
    FormatContext *fc = format_context;
    int offset = 0;
    IOContext *io = fc->io;
    long file_len;
    Stream *ost = fc->streams[0];

    pr_entry(in);

    file_len = io->io_get_len(io);
    io->io_seek(io, 0, SEEK_SET);

    fc->state = 1; /**< write packet */
    if (ost->track_type == 1) /**< if video */
        ovfp = fopen(fc->ofname, "w+");
    else
        oafp = fopen(fc->ofname, "w+");

    /* read all EBML elements. */
    do {
       ebml_parse(fc, mkv_syntax);
    } while (io->io_tell(io) < file_len);

    return 0;
}

Format format_mkv = {
    .formatname = "mkv",
    .id = ID_MKV,
    .read_probe = mkv_probe,
    .read_header = mkv_read_header,
    .read_packet = mkv_read_packet
};
